#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import rospy
import numpy as np
import tf
from sensor_msgs.msg import JointState
from geometry_msgs.msg import Pose, Point
from std_msgs.msg import Float64
import moveit_msgs.msg
import geometry_msgs.msg
import tf.transformations
import moveit_commander
from moveit_commander import MoveGroupCommander
from copy import deepcopy
import time

joint_state_data = None
result2_data = None
tf_viewer = None
allow_visual_control = True
class PIDController:
    def __init__(self, kp, ki, kd):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.prev_error = 0
        self.integral = 0
        self.last_time = time.time()

    def compute(self, setpoint, measured_value):
        current_time = time.time()
        dt = current_time - self.last_time
        if dt <= 0.0:
            return 0

        error = setpoint - measured_value
        self.integral += error * dt
        derivative = (error - self.prev_error) / dt

        output = self.kp * error + self.ki * self.integral + self.kd * derivative

        self.prev_error = error
        self.last_time = current_time

        return output


def joint_states_callback(data):
    global joint_state_data
    joint_state_data = data

def result2_callback(data):
    global result2_data
    result2_data = data

def process_data(event):
    global joint_state_data, result2_data, tf_viewer, allow_visual_control
    # if not allow_visual_control: 
    #     return
    if not allow_visual_control:
        raise KeyboardInterrupt

    if joint_state_data is not None and result2_data is not None:
            
            joint_positions = [joint_state_data.position[i] if i < len(joint_state_data.position) else 0 for i in range(5)]
            # 打印前五个关节的位置（弧度）
            # rospy.loginfo('joint_states: %s', joint_positions)
            # 打印前五个关节的位置(角度)
            joint_positions_degrees = [np.degrees(theta) for theta in joint_positions]
            joint_positions_degrees_rounded = [round(deg) for deg in joint_positions_degrees]
            # rospy.loginfo('joint_states (degrees): %s', joint_positions_degrees_rounded)

            # 打印/result2_topic接收到的xyz
            # rospy.loginfo('result2_topic xyz: %s', (result2_data.x, result2_data.y, result2_data.z))
            print("xyz:",result2_data.x,result2_data.y,result2_data.z)

            # 定义theta变量
            theta1, theta2, theta3, theta4, theta5 = joint_positions
            theta6 = 0

            cos = np.cos
            sin = np.sin

            TS = np.array([
                [cos(theta5)*sin(theta1) - sin(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))),
                 sin(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) - cos(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)), 
                 - sin(theta1)*sin(theta5) - cos(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))), 
                 181*cos(theta1)*sin(theta2) - (8693*sin(theta1)*sin(theta5))/100 + (77*cos(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)))/2 - (141*cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)))/2 - (141*sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2)))/2 - (77*sin(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)))/2 - (8693*cos(theta5)*(cos(theta4)*(cos(theta1)*sin(theta2)*sin(theta3) - cos(theta1)*cos(theta2)*cos(theta3)) + sin(theta4)*(cos(theta1)*cos(theta2)*sin(theta3) + cos(theta1)*cos(theta3)*sin(theta2))))/100 + 181*cos(theta1)*cos(theta2)*sin(theta3) + 181*cos(theta1)*cos(theta3)*sin(theta2) - 483/4],

                [- cos(theta1)*cos(theta5) - sin(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))), 
                 sin(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) - cos(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)), 
                 cos(theta1)*sin(theta5) - cos(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))), 
                 (8693*cos(theta1)*sin(theta5))/100 + 181*sin(theta1)*sin(theta2) + (77*cos(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)))/2 - (141*cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)))/2 - (141*sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2)))/2 - (77*sin(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)))/2 - (8693*cos(theta5)*(cos(theta4)*(sin(theta1)*sin(theta2)*sin(theta3) - cos(theta2)*cos(theta3)*sin(theta1)) + sin(theta4)*(cos(theta2)*sin(theta1)*sin(theta3) + cos(theta3)*sin(theta1)*sin(theta2))))/100 + 181*cos(theta2)*sin(theta1)*sin(theta3) + 181*cos(theta3)*sin(theta1)*sin(theta2)],

                [-sin(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))), 
                 sin(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) - cos(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)), 
                 -cos(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))), 
                 181*cos(theta2) + 181*cos(theta2)*cos(theta3) - 181*sin(theta2)*sin(theta3) - (8693*cos(theta5)*(cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)) + sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3))))/100 - (141*cos(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)))/2 + (77*cos(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)))/2 - (77*sin(theta4)*(cos(theta2)*sin(theta3) + cos(theta3)*sin(theta2)))/2 - (141*sin(theta4)*(cos(theta2)*cos(theta3) - sin(theta2)*sin(theta3)))/2 + 193/2],
                [0, 0, 0, 1]
            ])
            # rospy.loginfo('TS:\n %s', TS)

            global M
            M = np.eye(4) # 定义M矩阵
            M[0, 3] = -result2_data.x
            # 
            M[1, 3] = 0
            # result2_data.y
            M[2, 3] = result2_data.z
            # print(M[0, 3],M[1, 3],M[2, 3])

            Tc12 = np.array([[0, -1, 0, 0], [0, 0, -1, 0], [1, 0, 0, 0], [0, 0, 0, 1]])
            Tce = np.array([[1, 0, 0, 88.876], [0, 1, 0, 0], [0, 0, 1, -59.6856], [0, 0, 0, 1]])
            result_matrix = TS @ M @ Tc12# 结果矩阵计算 
            result_matrix[:3, 3] /= 1000  # 调整尺度
            # 对矩阵中的值进行四舍五入
            rounded_result_matrix = np.round(result_matrix, decimals=3)  # 四舍五入到小数点后 6 位
            # 打印四舍五入后的矩阵
            # rospy.loginfo('Rounded Result Matrix:\n %s', rounded_result_matrix)
            position = result_matrix[:3, 3] 
            rotation_matrix = result_matrix[:3, :3]
            # 打印位置和旋转矩阵
            print("Position:", position)
            # print("Rotation Matrix:\n", rotation_matrix)
            quaternion = tf.transformations.quaternion_from_matrix(result_matrix) # 将旋转矩阵转换为四元数
            # # 打印位置和四元数
            # print("Quaternion:\n", quaternion)

            move_arm_to_position(position, quaternion)
            # if success:
            #     print("规划成功,退出")
                # rospy.signal_shutdown("Shutting down node after successful motion plan")
                # print("111")

def move_arm_to_position(position, quaternion):
    global arm
    global M

    # 初始化 PID 控制器（针对 y 轴的控制）
    pid_y = PIDController(kp=0.002, ki=0.0003, kd=0.00)

    # 获取当前机械臂的位置
    current_pose = arm.get_current_pose().pose

    # 计算当前位置与目标位置的距离
    current_position = np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z])
    target_position = np.array(position)
    distance = np.linalg.norm(current_position - target_position)

    # 定义阈值（例如，0.01米）
    threshold = 0.07
    if distance < threshold:
        print("当前位置与目标位置相差不远，不做规划。")
        return  # 取消规划

    # 逐步逼近目标位置 1
    max_step1 = 0.02  # 定义每次移动的最大步长（例如，1厘米）
    step_vector1 = target_position[1] - current_position[1]
    step_distance1 = np.linalg.norm(step_vector1)

    if step_distance1 > max_step1:
        step_vector1 = (step_vector1 / step_distance1) * max_step1  # 计算按比例缩小的步长
    # 逐步逼近目标位置 2
    max_step2 = 1  # 定义每次移动的最大步长（例如，1厘米）
    step_vector2 = target_position[0] - current_position[0]
    step_distance2 = np.linalg.norm(step_vector2)

    if step_distance2> max_step2:
        step_vector2= (step_vector2/step_distance2) * max_step2  # 计算按比例缩小的步长

    # 使用 PID 控制调整 y 轴位置
    pid_output_y = pid_y.compute(target_position[1], current_position[1])
    intermediate_position = current_position
    intermediate_position[1] += step_vector1
    intermediate_position[0] += step_vector2
    # print(f"Current x position: {current_position[0]}, Target x position: {target_position[0]}")
    # print(f"Current y position: {current_position[1]}, Target y position: {target_position[1]}")
    # print(f"Current Z position: {current_position[2]}, Target Z position: {target_position[2]}")
    
    intermediate_position[1] += pid_output_y  # 调整 y 轴位置

    target_pose = Pose()  # 设置目标位姿
    target_pose.position.x = intermediate_position[0]
    target_pose.position.y = intermediate_position[1]
    target_pose.position.z = current_position[2]
    target_pose.orientation.x = quaternion[0]
    target_pose.orientation.y = quaternion[1]
    target_pose.orientation.z = quaternion[2]
    target_pose.orientation.w = quaternion[3]

    # 控制机械臂运动到目标位姿
    arm.set_pose_target(target_pose)
    success = arm.go()
    rospy.sleep(1)

    if success:
        if M[0, 3] == 0 and M[1, 3] == 0 and M[2, 3] == 0:
            print("保持不动\n")
        else:
            print("规划成功\n")
        return True
    else:
        print("未规划")
        return False
allow_visual_control = True  # 允许视觉

def stop_visual_control():
    global allow_visual_control
    allow_visual_control = False

def run_visual_control():
    global allow_visual_control
    allow_visual_control = True

def visual_main():
    global arm
    print("视觉开始")
    # rospy.init_node('multi_topic_listener', anonymous=True)  # 保留这个初始化节点的调用

    moveit_commander.roscpp_initialize(sys.argv) # 初始化move_group的API
    
    arm = MoveGroupCommander('manipulator') # 初始化需要使用move group控制的机械臂中的arm group

    arm.allow_replanning(True) # 当运动规划失败后，允许重新规划

    arm.set_pose_reference_frame('base_link') # 设置目标位置所使用的参考坐标系

    arm.set_goal_position_tolerance(0.0001) 
    arm.set_goal_orientation_tolerance(0.0001)

    arm.set_max_acceleration_scaling_factor(0.1)
    arm.set_max_velocity_scaling_factor(0.1)

    # 控制机械臂先回到初始化位置
    arm.set_named_target('middle')
    arm.go()
    rospy.sleep(1)

    global tf_viewer
    tf_viewer = rospy.Publisher('result_data', Float64, queue_size=10)
    rospy.Subscriber("/joint_states", JointState, joint_states_callback)
    rospy.Subscriber("/result3_topic", Point, result2_callback)

    # ## 模拟接收到 result2_topic 消息
    # global result2_data
    # result2_data = Point(0, 0, 200)

    rospy.Timer(rospy.Duration(2), process_data)
#     rospy.spin()


# if __name__ == '__main__':
#     try:
#         moveit_commander.roscpp_initialize(sys.argv)
#         visual_main()
#     except rospy.ROSInterruptException:
#         pass